import sys
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

sys.path.append('./')

import DataProcessing
import GFTA

day_date_r0v, day_open_r0v, day_high_r0v, day_low_r0v, day_close_r0v, day_volume_r0v = DataProcessing.data_processing.split_day_data()

week_date, week_open, week_high, week_low, week_close, week_volume = DataProcessing.data_processing.split_week_data()

day_ma5, day_ma10, day_ma20, day_ma30, day_ma60, day_ma120, day_ma250 = GFTA.ta_index.day_ma()

week_ma5, week_ma10, week_ma20, week_ma30, week_ma60, week_ma120, week_ma250 = GFTA.ta_index.week_ma()

day_dif, day_dea, day_macd = GFTA.ta_index.day_macd()

week_dif, week_dea, week_macd = GFTA.ta_index.week_macd()

class trade():

    def star():

        capital = 100000

        position = 0 #4000/2000/1000

        # 买入逻辑
        for wdate, wdif, wdea in zip(week_date, week_dif, week_dea):

            if position == 0:

                if wdif > wdea:

                    for ddate, dclose in zip(day_date_r0v, day_close_r0v):

                        if ddate == wdate:

                            position += 4000

                            capital -= dclose * 4000

                            print("日期:", ddate, "当前总资产:", capital + dclose * position, "可用资金:", capital, "当前持仓数量:", position, "当前持仓市值:", dclose * position)

class plot():

    def plot_setting():

        # 指定默认字体(解决中文无法显示的问题)
        plt.rcParams['font.sans-serif'] = ['SimHei']

        # 解决保存图像时负号“-”显示方块的问题
        plt.rcParams['axes.unicode_minus'] = False

    def plot_figure():

        # 创建一个8*6点的图，设置分辨率为72
        fig = plt.figure(figsize=(15,10), dpi=72)

    def plot_day():

        ax11 = plt.subplot2grid((5,2), (0,0), rowspan=3, colspan=1)
        
        # 绘制收盘价走势图(日)
        # --------------------------------------------------
        #ax11.plot(day_close_r0v, color='grey', alpha=0.3)
        #ax11.fill_between(range(0,len(day_close_r0v)), 0, day_close_r0v, facecolor='grey', alpha=0.3)

        # 绘制蜡烛图(日)
        # --------------------------------------------------
        for n, o, h, l, c in zip(range(0,len(day_close_r0v)), day_open_r0v, day_high_r0v, day_low_r0v, day_close_r0v):

            if c >= o:

                ax11.bar(n, (c - o), 0.8, bottom=o, color='white', edgecolor='red', zorder=3)

                ax11.vlines(n, l, h, color='red')

            else:

                ax11.bar(n, (c - o), 0.8, bottom=o, color='green', edgecolor='green', zorder=3)

                ax11.vlines(n, l, h, color='green')
        
        # 绘制MA均线(日)
        # --------------------------------------------------
        ax11.plot(day_ma5, label="MA 5", color='black')
        ax11.plot(day_ma10, label="MA 10", color='yellow')
        ax11.plot(day_ma20, label="MA 20", color='purple')
        ax11.plot(day_ma30, label="MA 30", color='pink')
        ax11.plot(day_ma60, label="MA 60", color='green')
        ax11.plot(day_ma120, label="MA 120", color='blue')
        ax11.plot(day_ma250, label="MA 250", color='brown')

        #ax11.xaxis.set_major_locator(ticker.MultipleLocator(500))
        #plt.xticks(day_date_r0v,rotation=45)

        ax11.set_title("收盘价走势图(日)")
        ax11.legend()
        #ax11.set_xlabel("日期")
        ax11.set_ylabel("收盘价(日)",rotation=90)

        # 绘制成交量(日)
        # --------------------------------------------------
        ax12 = plt.subplot2grid((5,2), (3,0), rowspan=1, colspan=1)
        ax12.bar(range(0,len(day_volume_r0v)), day_volume_r0v, color='red')
        ax12.set_title("成交量(日)")

        # 绘制MACD(日)
        # --------------------------------------------------
        ax13 = plt.subplot2grid((5,2), (4,0), rowspan=1, colspan=1)
        ax13.plot(day_dif, color='black')
        ax13.plot(day_dea, color='yellow')

        #ax13.bar(range(0,len(day_macd)), day_macd, color='red')
        for o, i in zip(range(0,len(day_macd)), day_macd):
            if i >= 0:
                ax13.bar(o, i, color='red')

            else:
                ax13.bar(o, i, color='green')

        ax13.set_title("MACD(日)")

    def plot_week():

        ax21 = plt.subplot2grid((5,2), (0,1), rowspan=3, colspan=1)
        
        # 绘制收盘价走势图(周)
        # --------------------------------------------------
        #ax21.plot(week_close, color='grey', alpha=0.3)
        #ax21.fill_between(range(0,len(week_close)), 0, week_close, facecolor='grey', alpha=0.3)

        # 绘制蜡烛图图(周)
        # --------------------------------------------------
        for n, o, h, l, c in zip(range(0,len(week_close)), week_open, week_high, week_low, week_close):

            if c >= o:

                ax21.bar(n, (c - o), 0.8, bottom=o, color='white', edgecolor='red', zorder=3)

                ax21.vlines(n, l, h, color='red')

            else:

                ax21.bar(n, (c - o), 0.8, bottom=o, color='green', edgecolor='green', zorder=3)

                ax21.vlines(n, l, h, color='green')
        
        # 绘制MA均线(周)
        # --------------------------------------------------
        ax21.plot(week_ma5, label="MA 5", color='black')
        ax21.plot(week_ma10, label="MA 10", color='yellow')
        ax21.plot(week_ma20, label="MA 20", color='purple')
        ax21.plot(week_ma30, label="MA 30", color='pink')
        ax21.plot(week_ma60, label="MA 60", color='green')
        ax21.plot(week_ma120, label="MA 120", color='blue')
        ax21.plot(week_ma250, label="MA 250", color='brown')
        ax21.set_title("收盘价走势图(周)")
        ax21.legend()
        ax21.set_ylabel("收盘价(周)",rotation=90)

        # 绘制成交量(周)
        # --------------------------------------------------
        ax22 = plt.subplot2grid((5,2), (3,1), rowspan=1, colspan=1)
        ax22.bar(range(0,len(week_volume)), week_volume, color='red')
        ax22.set_title("成交量(周)")

        # 绘制MACD(周)
        # --------------------------------------------------
        ax23 = plt.subplot2grid((5,2), (4,1), rowspan=1, colspan=1)
        ax23.plot(week_dif, color='black')
        ax23.plot(week_dea, color='yellow')
        
        #ax23.bar(range(0,len(week_macd)), week_macd, color='red')
        for o, i in zip(range(0,len(week_macd)), week_macd):
            if i >= 0:
                ax23.bar(o, i, color='red')

            else:
                ax23.bar(o, i, color='green')
        
        ax23.set_title("MACD(周)")

if __name__=='__main__':

    trade.star()

    plot.plot_setting()
    plot.plot_figure()
    plot.plot_day()
    plot.plot_week()

    plt.tight_layout()

    #plt.savefig('quxiantu.png',dpi=120,bbox_inches='tight')
    plt.show()
    #plt.close()

# - END -
